[Diffusion] Rollout API: support return only SDE latents - #36994

Motivation

/rollout/generate returns the whole 0..T denoising trajectory, but flowGRPO training consumes only the SDE window: for each SDE step s it needs x_s, x_{s+1} and log_prob_s. On a 10-step schedule with a 2-step window that is 3 of the 11 latents; the other 8 can be thrown away, cutting trajectory payload size down by ~72%. These latents are fp32 and, for video models, dominate the response body — a MiniMax H3 t2va microgroup carries hundreds of MB of them.

rollout_return_step_indices already existed for this, but it was not usable:

• the trainer could not tell which original step each surviving latent came from, so pairing x_s with x_{s+1} relied on the array being the full contiguous 0..T;
• timesteps was filtered together with the latents, so the trainer lost the schedule entries it indexes by step number;
• the MiniMax H3 path ignored the field entirely — its own collector recorded every step unconditionally.

Modifications

• RolloutDitTrajectory gains latent_step_indices: the original step index of each returned latent (None = the full 0..T trajectory). Serialized alongside the trajectory and carried through the per-sample split.
• RolloutDenoisingMixin: latents keep honoring rollout_return_step_indices and now record their provenance; timesteps is appended before the filter, so the schedule always comes back full-length. Scalar-sized arrays (timesteps, rollout_log_probs) are indexed by original step number, only the latents are filtered.
• MiniMaxH3RolloutCollector: same filtering and provenance, so H3 stops ignoring the request; its timesteps now spans the full [T+1] including the terminal entry, matching sigmas and the generic path.

The consumer side lives in radixark/miles_diffusion#216: it requests S ∪ (S+1) for an SDE window S and pairs latents by the echoed provenance instead of by array position, so a non-contiguous window is correct by construction and a request/response disagreement raises instead of silently mispairing.

Accuracy Tests

Joint end-to-end runs with radixark/miles_diffusion#216 pinned to this branch — all e2e stages green (run 33244938347): stage-b-3-gpu-h200, stage-b-5-gpu-h200, stage-c-3-gpu-h200, stage-c-5-gpu-h200.

Every one of those recipes runs a step strategy, so all of them exercised the filtered transport, and each reproduced its pre-existing recorded metric standard bit for bit (strict comparison, --deterministic-mode):

Two further e2e CIs are being added for the models this PR touches. Each recorded its standard off this branch, then reproduced it on top of it:

cosmos3 is the case that cannot pass without the provenance echo: its Karras grid puts the useful window at steps 8-11, so array position never coincides with step number and pairing by position would silently pair x_10 with x_0. Its train/model_output_{mean_abs_diff,rel_max} compare the raw DiT outputs between engine and trainer, so the equality covers the tensors themselves rather than a reward summary.

Checklist

• Format your code according to the Format code with pre-commit. — run on the touched files; black reformatting is included in the commit
• Add unit tests according to the Run and add unit tests. — no new unit test; the filter and provenance were verified behaviourally against the H3 collector and the generic mixin, and by the bitwise e2e comparisons above
• Update documentation according to Write documentations. — not needed, the field is described where it is declared
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ❌ Run #33282596522
Latest PR Test (Extra): ❌ Run #33282596363
Latest PR Test (AMD ROCm 7.2): ❌ Run #33282596623